import multiprocessing
import os
import sys
from urllib.request import urlopen, urlretrieve

import six


def Worker(url_and_target):
    try:
        (url, target) = url_and_target
        print('>>> downloading ' + target)
        urlretrieve(url, target)
        # if not os.path.isfile(target):
        #     print('>>> downloading ' + target)
        #     urlretrieve(url, target)
        # else:
        #     print(f'>>> downloaded {target}')
    except (KeyboardInterrupt, SystemExit):
        print('>>> exit child process')


class KGSIndex:
    def __init__(self, kgs_url=r'http://u-go.net/gamerecords/',
                 index_page='kgs_index.html', data_dir='kgs') -> None:
        self.kgs_url = kgs_url
        self.index_page=index_page
        self.data_dir = data_dir
        self.file_info = []
        self.urls = []
        self.LoadIndex()

    def DownloadFiles(self):
        if not os.path.isdir(self.data_dir):
            os.makedirs(self.data_dir)
        urls = []
        for file_info in self.file_info:
            url = file_info['url']
            filename = file_info['filename']
            if not os.path.isfile(self.data_dir + '/' + filename):
                urls.append((url, self.data_dir + '/' + filename))
        cores = multiprocessing.cpu_count()
        pool = multiprocessing.Pool(processes=cores)
        try:
            it = pool.imap(Worker, urls)
            for _ in it:
                pass 
            pool.close()
            pool.join()
        except KeyboardInterrupt:
            print('>>> Caught KeyboardInterrupt, terminating worker')
            pool.terminate()
            pool.join()
            sys.exit(-1)

    def LoadIndex(self):
        index_contents = self.CreateIndexPage()
        split_page = [item for item in index_contents.split('<a href="') if item.startswith('https://')]
        for item in split_page:
            download_url = item.split('">Download')[0]
            if download_url.endswith('.tar.gz'):
                self.urls.append(download_url)
        for url in self.urls:
            filename = os.path.basename(url)
            split_filename = filename.split('-')
            num_games = int(split_filename[len(split_filename)-2])
            print(f'{filename} {num_games}')
            self.file_info.append({'url': url, 'filename': filename, 'num_games': num_games})

    def CreateIndexPage(self): 
        if os.path.isfile(self.index_page):
            print('>>> Reading cached index page.')
            index_file = open(self.index_page, 'r')
            index_contents = index_file.read()
            index_file.close()
        else:
            print('>>> Downloading index page.')
            fp = urlopen(self.kgs_url)
            data = six.text_type(fp.read())
            fp.close()
            index_contents = data 
            index_file = open(self.index_page, 'w')
            index_file.write(index_contents)
            index_file.close()
        return index_contents
    
def run():
    index = KGSIndex()
    index.DownloadFiles()

if __name__ == '__main__':
    run()